(import
 (except (rnrs base)
         let-values
         map
         error
         vector-map)
 (only (guile)
       lambda* λ
       simple-format
       current-output-port)
 (fileio)
 (queue)
 (list-helpers)
 (math)
 ;; lists
 (srfi srfi-1)
 ;; vectors
 (srfi srfi-43)
 ;; let-values
 ;; (srfi srfi-11)
 ;; hash tables
 ;; (srfi srfi-69)
 ;; functional records
 (srfi srfi-9 gnu)
 (ice-9 pretty-print)
 (ice-9 peg)
 (prefix (peg-tree-utils) peg-tree:))

;; define parser

(define-peg-pattern SPACE none " ")
(define-peg-pattern NEWLINE none "\n")
(define-peg-pattern COLON none ":")
(define-peg-pattern KEY-VAL-SEP none (and COLON SPACE))
(define-peg-pattern NUMBER body (+ (range #\0 #\9)))
(define-peg-pattern ANYTHING-EXCEPT-NUMBER none
  (* (and (not-followed-by NUMBER) peg-any)))

;; 1. line
(define-peg-pattern MONKEY-LABEL none "Monkey")
(define-peg-pattern ID all NUMBER)

;; 2. line
(define-peg-pattern NUMBER-LIST all (* (and NUMBER (? ",") (? SPACE))))
(define-peg-pattern ITEMS-LINE all (and ANYTHING-EXCEPT-NUMBER NUMBER-LIST))

;; 3. line
(define-peg-pattern OP-LABEL none "Operation")
(define-peg-pattern OP-VAR-NEW none "new")
(define-peg-pattern OP-EQUALS none "=")
(define-peg-pattern OP-VAR-OLD none "old")

(define-peg-pattern OPERATOR all (or "+" "*"))
(define-peg-pattern OPERAND-OLD body "old")
(define-peg-pattern OPERAND all (or NUMBER OPERAND-OLD))

(define-peg-pattern OPERATION all
  (and OPERATOR SPACE OPERAND))

(define-peg-pattern ID-LINE all
  (and MONKEY-LABEL SPACE ID COLON))

(define-peg-pattern OP-LINE all
  (and (* SPACE) OP-LABEL KEY-VAL-SEP
       OP-VAR-NEW SPACE OP-EQUALS SPACE OP-VAR-OLD SPACE
       OPERATION))

;; 4. line
(define-peg-pattern TEST-DIVISOR all NUMBER)
(define-peg-pattern TEST-LINE all (and ANYTHING-EXCEPT-NUMBER TEST-DIVISOR))

;; 5. line
(define-peg-pattern TEST-TRUE-NEXT-MONKEY all NUMBER)
(define-peg-pattern TEST-TRUE-LINE all (and ANYTHING-EXCEPT-NUMBER TEST-TRUE-NEXT-MONKEY))

;; 6. line
(define-peg-pattern TEST-FALSE-NEXT-MONKEY all NUMBER)
(define-peg-pattern TEST-FALSE-LINE all (and ANYTHING-EXCEPT-NUMBER TEST-FALSE-NEXT-MONKEY))

(define-peg-pattern MONKEY all
  (and ID-LINE NEWLINE
       ITEMS-LINE NEWLINE
       OP-LINE NEWLINE
       TEST-LINE NEWLINE
       TEST-TRUE-LINE NEWLINE
       TEST-FALSE-LINE))

;; define model

(define-immutable-record-type <monkey>
  (make-monkey id items inspections operation test true-next false-next)
  monkey?
  (id monkey-id)
  (operation monkey-operation)
  (test monkey-test)
  (true-next monkey-true-next)
  (false-next monkey-false-next)
  (items monkey-items set-monkey-items)
  (inspections monkey-inspections set-monkey-inspections))


(set-record-type-printer!
 <monkey>
 (lambda (record port)
   (simple-format port
                  "<monkey: id:~a items:~a inspections:~a true-next:~a false-next:~a>"
                  (monkey-id record)
                  (monkey-items record)
                  (monkey-inspections record)
                  (monkey-true-next record)
                  (monkey-false-next record))))


(define-syntax ->
  (syntax-rules ()
    ;; first expression is left unchanged
    [(-> expr) expr]
    ;; take from the back, wrap other calls
    [(-> expr* ... (op args* ...))
     (op args* ... (-> expr* ...))]
    ;; make parens unnecessary in trivial case of no further arguments
    [(-> expr* ... op)
     (op (-> expr* ...))]))


(define lines (get-lines-from-file "input"))


(define monkey-lines
  (split-into-segments lines (λ (line) (string-null? line))))


(define string-operator->operator
  (λ (str)
    (cond
     [(string=? str "+") +]
     [(string=? str "*") *]
     [else (error "unrecognized operator" str)])))


(define monkey-lines->monkey
  (λ (lines)
    (let* ([monkey-str (string-join lines "\n")]
           [tree (peg:tree (match-pattern MONKEY monkey-str))])
      (let ([id (string->number (car (peg-tree:tree-refs tree '(ID))))]
            [items
             (map string->number
                  (string-split (car (peg-tree:tree-refs tree '(ITEMS-LINE NUMBER-LIST)))
                                #\,))]
            [inspections 0]
            ;; Operation is a function taking the previous
            ;; worry value as argument.
            [operation
             (let ([operand-str (car (peg-tree:tree-refs tree '(OPERATION OPERAND)))]
                   [operator
                    (string-operator->operator
                     (car (peg-tree:tree-refs tree '(OPERATION OPERATOR))))])
               ;; Handling the annoying special case ...
               (cond
                [(string=? operand-str "old")
                 (λ (worry)
                   (simple-format (current-output-port)
                                  "handling item of worry: ~a\n"
                                  worry)
                   (operator worry worry))]
                [else
                 (λ (worry)
                   (simple-format (current-output-port)
                                  "handling item of worry: ~a\n"
                                  worry)
                   (operator worry (string->number operand-str)))]))]
            [test-divisor
             (string->number
              (car (peg-tree:tree-refs tree '(TEST-DIVISOR))))]
            [test-true-next-monkey
             (string->number
              (car (peg-tree:tree-refs tree '(TEST-TRUE-NEXT-MONKEY))))]
            [test-false-next-monkey
             (string->number
              (car (peg-tree:tree-refs tree '(TEST-FALSE-NEXT-MONKEY))))])
        ;; (make-monkey id items inspections operation test true-next false-next)
        (make-monkey id
                     items
                     0  ; inspections so far
                     operation
                     (λ (worry)
                       (simple-format (current-output-port)
                                      "~a divisible-by? ~a -> ~a\n"
                                      worry test-divisor
                                      (divisible-by? worry test-divisor))
                       (divisible-by? worry test-divisor))
                     test-true-next-monkey
                     test-false-next-monkey)))))


(define monkeys
  (list->vector
   (map monkey-lines->monkey
        monkey-lines)))


(define bored
  (λ (item-worry)
    (-> item-worry
        ((λ (worry) (/ worry 3)))
        floor)))


(define move-items
  (λ (monkeys ind)
    (let iter ([items° (monkey-items (vector-ref monkeys ind))]
               [monkeys° monkeys])
      (let* ([monkey (vector-ref monkeys ind)]
             [monkey-op (monkey-operation monkey)])
        (cond
         [(null? items°) monkeys°]
         [else
          (let* ([new-item-worry (-> items° car monkey-op bored)])
            (let ([target-monkey-ind
                   (if ((monkey-test monkey) new-item-worry)
                       (monkey-true-next monkey)
                       (monkey-false-next monkey))])
              (simple-format (current-output-port)
                             "moving item; was worry ~a now worry ~a from index: ~a to index: ~a\n"
                             (car items°) new-item-worry ind target-monkey-ind)
              ;; Remove item from current monkey's item list
              ;; and add to inspection counter.
              (vector-set! monkeys°
                           ind
                           (set-fields monkey
                                       ((monkey-inspections)
                                        (+ (monkey-inspections (vector-ref monkeys° ind)) 1))
                                       ((monkey-items) (cdr items°))))
              ;; Append the item to the target monkey's item
              ;; list.
              (let* ([target-monkey (vector-ref monkeys° target-monkey-ind)]
                     [target-monkey-items
                      (append (monkey-items target-monkey)
                              (list new-item-worry))]
                     [updated
                      (set-monkey-items target-monkey target-monkey-items)])
                (vector-set! monkeys° target-monkey-ind updated))
              (iter (cdr items°) monkeys°)))])))))


(define calc-rounds
  (λ (monkeys num-rounds)
    (let iter ([round° 0] [monkeys° monkeys] [monkey-ind 0])
      (simple-format (current-output-port) "state:\n")
      (pretty-print monkeys°)
      (cond
       [(>= round° num-rounds)
        (simple-format (current-output-port) "done simulating rounds\n")
        monkeys°]
       [(>= monkey-ind (vector-length monkeys°))
        (simple-format (current-output-port) "simulate round ~a\n" (+ round° 1))
        (iter (+ round° 1) monkeys° 0)]
       [else
        (simple-format (current-output-port) "moving items of monkey ~a\n" monkey-ind)
        (iter round°
              (move-items monkeys° monkey-ind)
              (+ monkey-ind 1))]))))


(define monkey-business
  (λ (monkeys)
    (product
     (n-highest
      (vector->list
       (vector-map (λ (ind monkey) (monkey-inspections monkey))
                   monkeys))
      2))))


(define after-20-rounds (calc-rounds monkeys 20))
(simple-format (current-output-port) "monkeys after 20 rounds:\n")
(pretty-print after-20-rounds)
(define result (monkey-business after-20-rounds))
(simple-format (current-output-port)
               "~a\n"
               (vector->list
                (vector-map (λ (ind monkey) (monkey-inspections monkey))
                            monkeys)))
(simple-format (current-output-port) "result: ~a\n" result)
